# -*- coding: utf-8 -*-
import os
import sys
import argparse
import random
import csv


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--dataset", help="Folder with data", required=True)
    parser.add_argument("-n", "--test_size", help="Number of test samples per folder. Default value is 2", required=False, type=float, default=0.1)
    parser.add_argument("--no-count", help="Do not write number of samples at first line of files", required=False, action='store_true')
    parser.add_argument("--class-threshold", help="Threshold for number of classes. All for default", required=False,  type=int, default=0)
    return parser.parse_args()
    

if __name__ == '__main__':
    args = parse_args()
    
    test_size = args.test_size
    
    class_id = 0
    number_test_samples = 0
    number_train_samples = 0
    
    test_file = open('test.csv', 'wb')
    train_file = open('train.csv', 'wb')
    if not args.no_count:
        test_file.write('%05d\n' % number_test_samples)  # Hack for creating space in the beginning of file for inserting real number when work will be done
        train_file.write('%05d\n' % number_train_samples)
    test_writer = csv.writer(test_file, delimiter=';')
    train_writer = csv.writer(train_file, delimiter=';')
    test_writer.writerow(['File path', 'Real class', 'Class id'])
    train_writer.writerow(['File path', 'Real class', 'Class id'])
    
    for root, dirs, files in os.walk(args.dataset):
        filenames = [os.path.join(root, name) for name in files]
        folder_name = os.path.split(root)[-1].strip(' .0123456789')  #Clear digits and leading\trailing spaces
        filenames = [x for x in filenames if x.split('.')[-1].lower() in ['png', 'tiff', 'jpg']]
        if not filenames:
            continue
        elif len(filenames) == 1:
            print '%s: Only one sample in set, impossible to separate, will skip it' % folder_name
            continue
            
        class_id += 1
        if args.class_threshold > 0 and class_id > args.class_threshold:
            print 'Class threshold reached. Finish..'
            break

        try:
            test_samples = random.sample(filenames, int(round(test_size * len(filenames))))
        except ValueError as exc:
            print '%s: Specified test_size too big, there are only %s samples' % (folder_name, len(filenames))
            print '\tWill use 1 test sample'
            test_samples = random.sample(filenames, 1)
        train_samples = [x for x in filenames if x not in test_samples]
        
        number_test_samples += len(test_samples)
        number_train_samples += len(train_samples)
        for sample in test_samples:
            test_writer.writerow([sample, folder_name, class_id])
        for sample in train_samples:
            train_writer.writerow([sample, folder_name, class_id])
        
    if not args.no_count:
        test_file.seek(0)
        train_file.seek(0)
        test_file.write('%05d\n' % number_test_samples) 
        train_file.write('%05d\n' % number_train_samples)
    test_file.close()
    train_file.close()
